//
//  IORNNTrainer.cpp
//  iornn-parser-c
//
//  Created by Phong Le on 26/12/14.
//  Copyright (c) 2014 Phong Le. All rights reserved.
//

#include "IORNNTrainer.h"
#include "Reranker.h"

void IORNNTrainer::adagrad(IORNN* net, Param* grad, unordered_set<int>& tWords, unordered_set<int>& tPath) {
    real paramClr = paramLearningRate / (1 + nEvals*learningRateDecay);
    real wembClr = wembLearningRate / (1 + nEvals*learningRateDecay);
    
    if (paramVarience == NULL) {
        paramVarience = net->createGrad();
    }
    
    // for weights # Wword, L
    for (int i = 0; i < grad->nMatricesWoWE; i++) {
        Matrix* W = grad->weights[i];
        Matrix* temp = W->mul(W);
        paramVarience->weights[i]->addi(temp); delete temp;
        Matrix* std = paramVarience->weights[i]->sqrt()->addi(1e-10f);
        net->params->weights[i]->addi(grad->weights[i]->divi(std)->muli(-paramClr));
        delete std;
    }
    
    // for Wword
    for (unordered_set<int>::iterator it = tPath.begin(); it != tPath.end(); it++) {
        int p = *it;
        Matrix* temp = grad->Wword[p]->mul(grad->Wword[p]);
        paramVarience->Wword[p]->addi(temp); delete temp;
        Matrix* std = paramVarience->Wword[p]->sqrt()->addi(1e-10f);
        net->params->Wword[p]->addi(grad->Wword[p]->divi(std)->muli(-paramClr));
        delete std;
    }
    
    // for L
    for (unordered_set<int>::iterator it = tWords.begin(); it != tWords.end(); it++) {
        int word = *it;
        Matrix* temp = grad->L[word]->mul(grad->L[word]);
        paramVarience->L[word]->addi(temp); delete temp;
        Matrix* std = paramVarience->L[word]->sqrt()->addi(1e-10f);
        net->params->L[word]->addi(grad->L[word]->divi(std)->muli(-wembClr));
        delete std;
    }
    
    nEvals++;
}

void IORNNTrainer::train(IORNN* net, Treebank* tb, Treebank* devTb, KBestTreebank* kDevTb) {
    int nSample = tb->size();
    
    Param** gradList = new Param*[nThreads];
    for (int i = 0; i < nThreads; i++)
        gradList[i] = net->createGrad();
    
    int epoch = 0;
    struct timeval start, finish;
    
    int j = 0;
    real percent = 0;
    real percentS = 0;
    Reranker reranker;
    
    real* bestScore = NULL;
    int decreaseCount = 0;
    
    
    epoch++;
    cout << "=========== epoch " << epoch << " ========" << endl;;
    cout << Utils::currentDateTime() << endl;
    
    while (true) {
        j++;
        int startId = (j-1) * batchSize;
        int endId = min(nSample-1, j*batchSize-1);
        
        if (startId >= nSample) {
            if (epoch % evalDevStep == 0) {
                net->save(modelDir + "/model" + std::to_string(epoch));
            
                // eval
                cout << "++++++ EVAL... ++++++" << endl;
                gettimeofday(&start, NULL);
                real* as = reranker.optimizeMixReranker(net, devTb, kDevTb);
                gettimeofday(&finish, NULL);
                printf("[%.2fs]\n", ((double)(finish.tv_sec-start.tv_sec)*1000000 + (double)(finish.tv_usec-start.tv_usec)) / 1000000);
                
                if (bestScore == NULL) {
                    bestScore = as;
                    decreaseCount = 0;
                }
                else if (bestScore[0] <= as[0]) {
                    delete[] bestScore;
                    bestScore = as;
                    decreaseCount = 0;
                }
                
                else {
                    delete[] as;
                    decreaseCount++;
                }
                
                cout << decreaseCount << " consecutive decreases" << endl;
                
                // early stop
                if (decreaseCount >= maxDecreaseCount)
                    break;
            }
            
            j = 1;
            startId = (j-1) * batchSize;
            endId = min(nSample-1, j*batchSize-1);
            epoch++;
            percentS = 0;
            
            if (epoch > maxNEpoch) break;
            cout << "=========== epoch " << epoch << " ========" << endl;;
            cout << Utils::currentDateTime() << endl;
        }
        
        gettimeofday(&start, NULL);
        void** ret = net->computeCostAndGrad(tb, startId, endId, gradList);
        
        real cost = *(real*)ret[0];
        unordered_set<int> *tWords = (unordered_set<int>*)ret[1];
        unordered_set<int> *tPath = (unordered_set<int>*)ret[2];
    	adagrad(net, gradList[0], *tWords, *tPath);
        
        if (j % 100 == 0) {
        	cout << "batch " << j << "\t" << cost << "\t";
        	gettimeofday(&finish, NULL);
        	printf("[%.2fs]\n", ((double)(finish.tv_sec-start.tv_sec)*1000000 + (double)(finish.tv_usec-start.tv_usec)) / 1000000);
        }
        
        percent = endId * 100. / nSample;
        if (percent >= percentS) {
            cout << Utils::currentDateTime() << "\t";
            printf("%.2f%%", percent);
            percentS += 20;
            cout << endl;
        }
        
        delete[] ret;
        delete tWords;
        delete tPath;
    }
}



